import pandas as pd
import statsmodels.api as sm
from statsmodels.stats.outliers_influence import variance_inflation_factor
from statsmodels.tools.tools import add_constant


def do_analysis(df: pd.DataFrame, var_indexes: [], d_var_indexes: []):
    """
    计算回归分析
    :param df:
    :param var_indexes: 自变量
    :param d_var_indexes: 因变量
    :return:
    """

    def get_result(obj: pd.DataFrame) -> []:
        """
        获取前端需要的结结构
        :param obj:
        :return:
        """
        output = []
        for index, row in obj.iterrows():
            output.append({'row_key': str(index), 'values': {str(col): row[col] for col in obj.columns}})
        return output

    if len(var_indexes) <= 0 or len(d_var_indexes) <= 0:
        return
    cs_name = 'middle_calculate_column_diff_with_other'  # 初始化常量
    df[cs_name] = 1
    x_arr = [cs_name]
    for i in var_indexes:
        x_arr.append(i)
    x = df[x_arr]
    # y = pd.Series
    # if len(d_var_indexes) == 1:
    #     y = df[d_var_indexes[0]]
    # else:
    #     y = df[d_var_indexes]
    y = df[d_var_indexes]
    result = sm.OLS(y, x).fit()

    rf = result.rsquared  # R方
    rf_adj = result.rsquared_adj  # 调整后R方

    ess = result.ess  # 平方和
    df_model = result.df_model  # 自由度
    fvalue = result.fvalue  # f值 - F
    pvalue = result.f_pvalue  # p值 - 显著性
    mse_model = result.mse_model  # 均方

    # 系数
    params = result.params  # B
    df_b = pd.DataFrame({"B": params})
    bse = result.bse  # 标准错误
    df_bse = pd.DataFrame({"标准错误": bse})
    tvalues = result.tvalues  # t
    df_tvalues = pd.DataFrame({"t": tvalues})
    pvalues = result.pvalues  # p
    df_pvalues = pd.DataFrame({"显著性": pvalues})
    res_df = df_b.join(df_bse).join(df_tvalues).join(df_pvalues)  # 合并系数的每一列数据
    coefficients = get_result(res_df)  # 转换为输出结构
    res = {'summary': {'rf': rf, 'rf_adj': rf_adj, 'ess': ess, 'df_model': df_model, 'fvalue': fvalue, 'pvalue': pvalue,
                       'mse_model': mse_model},
           'coefficients': coefficients
           }
    return res


class RegressiveAnalysis(object):
    """
    回归分析
    """

    def __init__(self, df, var_indexes, d_var_indexes, regressive_r):
        self.df = df
        self.var_indexes = var_indexes
        self.d_var_indexes = d_var_indexes
        self.regressive_r = regressive_r
        self.vif_dict: dict = {}

    def get_result(self, obj: pd.DataFrame) -> []:
        """
        获取前端需要的结结构
        :param obj:
        :return:
        """
        output = []
        for index, row in obj.iterrows():
            row_key = str(index)
            values_dict = {str(col): row[col] for col in obj.columns}
            if self.vif_dict:
                values_dict['vif'] = self.vif_dict.get(row_key)
            output.append({'row_key': row_key, 'values': values_dict})
        return output

    def set_vif_dict(self, inner_df: pd.DataFrame):
        # 假设你有一个名为df的DataFrame，其中包含自变量
        # 为了计算VIF值，需要将自变量矩阵添加常数列
        cur_x = add_constant(inner_df)
        # 计算VIF值
        vif = pd.DataFrame()
        vif["Variable"] = cur_x.columns
        vif["VIF"] = [variance_inflation_factor(cur_x.values, i) for i in range(cur_x.shape[1])]
        res = dict()
        for index, row in vif.iterrows():
            res[str(row["Variable"])] = row["VIF"]
        # {'middle_calculate_column_diff_with_other': 16.526750331862086, 2: 1.4556519086023416,
        # 3: 1.3572137121310208, 4: 1.4549326597563321}
        self.vif_dict = res

    def analysis(self):
        """
        处理回归分析
        :return:
        """
        score = 0
        r_sq_pass = False  # R方是否通过 - 有一个通过则通过
        # p_pass_count = 0  # 显著性是否通过_数量计算
        # p_pass = False  # 显著性是否通过
        # vif_pass_count = 0  # vif是否通过
        var_indexes = self.var_indexes
        d_var_indexes = self.d_var_indexes
        df = self.df
        if len(var_indexes) <= 0 or len(d_var_indexes) <= 0:
            return
        cs_name = '常量'  # 初始化常量
        df[cs_name] = 1
        x_arr = [cs_name]
        for i in var_indexes:
            x_arr.append(i)
        x = df[x_arr]
        self.set_vif_dict(x)  # 计算vif
        res = []
        item_group_pass = False  # 用来判断是否有一组数据通过
        for d_var in d_var_indexes:
            y = df[[d_var]]
            result = sm.OLS(y, x).fit()

            rf = result.rsquared  # R方
            rf_adj = result.rsquared_adj  # 调整后R方
            rf_adj_f = float(rf_adj)
            score += rf_adj_f * 20
            if rf_adj_f >= self.regressive_r:  # 如果大于给定值就通过
                r_sq_pass = True  # 通过

            ess = result.ess  # 平方和
            df_model = result.df_model  # 自由度
            fvalue = result.fvalue  # f值 - F
            pvalue = result.f_pvalue  # p值 - 显著性
            mse_model = result.mse_model  # 均方

            # 系数
            params = result.params  # B
            df_b = pd.DataFrame({"B": params})
            bse = result.bse  # 标准错误
            df_bse = pd.DataFrame({"标准错误": bse})
            tvalues = result.tvalues  # t
            df_tvalues = pd.DataFrame({"t": tvalues})
            pvalues = result.pvalues  # p
            df_pvalues = pd.DataFrame({"显著性": pvalues})
            p_arr = pvalues.values
            p_pass_count = 0  # 显著性是否通过_数量计算
            for item_p_index in range(len(p_arr)):
                item_p = p_arr[item_p_index]
                if item_p_index == 0:  # 常数直接跳过
                    continue
                if item_p <= 0.05:
                    score += 10
                    p_pass_count += 1
            inner_p_pass = p_pass_count == len(var_indexes)  # 所有p值是否通过（p<=0.05）
            if not item_group_pass:
                item_group_pass = r_sq_pass and inner_p_pass
            res_df = df_b.join(df_bse).join(df_tvalues).join(df_pvalues)  # 合并系数的每一列数据
            coefficients = self.get_result(res_df)  # 转换为输出结构
            res.append({'summary': {'rf': rf, 'rf_adj': rf_adj, 'ess': ess, 'df_model': df_model, 'fvalue': fvalue,
                                    'pvalue': pvalue, 'mse_model': mse_model}, 'coefficients': coefficients})
        # print('var_indexes', var_indexes)
        # print('d_var_indexes', d_var_indexes)
        # print('p_pass_count', p_pass_count)
        vif_values = list(self.vif_dict.values())
        vif_pass_count = 0  # 用于判断vif是否通过
        for item_vif_index in range(len(vif_values)):
            item_vif = vif_values[item_vif_index]
            if item_vif_index == 0:  # 常数直接跳过
                continue
            score -= item_vif  # vif越小越好
            if item_vif <= 5:
                vif_pass_count += 1
        vif_pass = vif_pass_count == len(vif_values) - 1  # 所有vif值是否小于5
        # is_pass = r_sq_pass and p_pass and vif_pass
        # print('r_sq_pass:', r_sq_pass, 'p_pass:', p_pass, 'vif_pass', vif_pass)
        is_pass = item_group_pass and vif_pass
        return res, is_pass, score

    def report_analysis(self):
        """
        处理回归分析
        :return:
        """
        var_indexes = self.var_indexes
        d_var_indexes = self.d_var_indexes
        df = self.df
        if len(var_indexes) <= 0 or len(d_var_indexes) <= 0:
            return
        cs_name = '常量'  # 初始化常量
        df[cs_name] = 1
        x_arr = [cs_name]
        for i in var_indexes:
            x_arr.append(i)
        x = df[x_arr]
        self.set_vif_dict(x)  # 计算vif
        res = []
        for d_var in d_var_indexes:
            y = df[[d_var]]
            result = sm.OLS(y, x).fit()

            rf = result.rsquared  # R方
            rf_adj = result.rsquared_adj  # 调整后R方

            ess = result.ess  # 平方和
            df_model = result.df_model  # 自由度
            fvalue = result.fvalue  # f值 - F
            pvalue = result.f_pvalue  # p值 - 显著性
            mse_model = result.mse_model  # 均方

            # 系数
            params = result.params  # B
            df_b = pd.DataFrame({"B": params})
            bse = result.bse  # 标准错误
            df_bse = pd.DataFrame({"标准错误": bse})
            tvalues = result.tvalues  # t
            df_tvalues = pd.DataFrame({"t": tvalues})
            pvalues = result.pvalues  # p
            df_pvalues = pd.DataFrame({"显著性": pvalues})
            p_arr = pvalues.values
            p_pass_count = 0  # 显著性是否通过_数量计算
            for item_p_index in range(len(p_arr)):
                item_p = p_arr[item_p_index]
                if item_p_index == 0:  # 常数直接跳过
                    continue
                if item_p <= 0.05:
                    p_pass_count += 1
            res_df = df_b.join(df_bse).join(df_tvalues).join(df_pvalues)  # 合并系数的每一列数据
            coefficients = self.get_result(res_df)  # 转换为输出结构
            res.append({'summary': {'rf': rf, 'rf_adj': rf_adj, 'ess': ess, 'df_model': df_model, 'fvalue': fvalue,
                                    'pvalue': pvalue, 'mse_model': mse_model},
                        'coefficients': coefficients,
                        'y': d_var,
                        'x_arr': var_indexes})
        return res
